from tree_sitter import Language, Parser, Tree, Node
import tree_sitter_cpp
import os, json
import random
import string
from .type_get import TypeGet
from .identifier_get import IdentifierGet


class FileAnalyze:
    """
    methods:
    1. parse: parse the code and generate the abstract syntax tree
    2. analyze_function_definition: analyze the function definition and get the function name, start point, end point, parameters and function calls
    3. analyze_function_declarations: analyze the function declarations and get the function name, start point, end point and parameters
    4. classify_function_calls: classify the function calls into defined, declared only and undefined
    5. analyze_global_and_external_variables: analyze the global and external variables and get the variable name, start point and end point
    6. analyze_function_parameters: analyze the function parameters and get the type and declarator
    """
    def __init__(self, code: bytes, defined_types_input: dict = None):
        self.tree: Tree = self.parse(code)
        self.function_definitions: dict = {}
        self.function_declarations: dict = {}
        self.analyze_function_definition()
        self.analyze_function_declarations()
        self.classified_functions: dict = {}
        self.classify_function_calls()
        self.defined_types: dict = {}
        # self.analyze_struct_definitions()
        if defined_types_input:
            self.defined_types = defined_types_input
        else:
            with open("/disk0/wyf/Project/zz_static_project/type_definitions.json", "r") as file:
                self.defined_types = json.load(file)
        self.global_variables = {}
        self.external_variables = {}
        self.analyze_global_and_external_variables()

        self.analyze_function_transferability()

    def parse(self, code: bytes):
        CPP_LANGUATE = Language(tree_sitter_cpp.language())
        parser = Parser(CPP_LANGUATE)
        try:
            return parser.parse(code)
        except Exception as e:
            print(f"Failed to parse ast with exception: {e}")
            return
        
    def analyze_function_definition(self) -> None:
        def traverse(node: Node, parent_function=None):
            if node.type == "function_definition":
                function_name = None
                function_return_type = None
                function_parameters = []
                local_variables = {}
                identifiers = []
                exception_flag = False
                for child in node.children:
                    if child.type in type_specifier:
                        function_return_type = self.analyze_function_return_type(child)
                        if function_return_type is False:
                            exception_flag = True
                        continue
                    elif child.type == "function_declarator":
                        for grandchild in child.children:
                            if grandchild.type == "identifier":
                                function_name = grandchild.text.decode("utf-8")
                                continue
                            if grandchild.type == "parameter_list":
                                function_parameters = self.analyze_function_parameters(grandchild)
                                if function_parameters is False:
                                    exception_flag = True
                                continue
                        continue
                    elif child.type == "compound_statement":
                        local_variables = self.analyze_local_variables(child)
                        identifiers = self.analyze_identifiers(child)
                        if local_variables is False or identifiers is False:
                            exception_flag = True
                        continue
                if function_name:
                    function_start = node.start_point
                    function_end = node.end_point                    
                    self.function_definitions[function_name] = {
                        "name": function_name,
                        "start": function_start,
                        "end": function_end,
                        "return_type": function_return_type,
                        "parameters": function_parameters,
                        "local_variables": local_variables,
                        "identifiers": identifiers,
                        "calls": set(),
                        "exception_flag": exception_flag,
                    }
                    parent_function = function_name

            elif node.type == "call_expression" and parent_function:
                for child in node.children:
                    if child.type == "identifier":
                        called_function = child.text.decode("utf-8")
                        self.function_definitions[parent_function]["calls"].add(called_function)
                        break

            for child in node.children:
                traverse(child, parent_function)

        if self.tree and self.tree.root_node:
            typeget = TypeGet()
            type_specifier = typeget.get_subtype_of_type_specifier()
            traverse(self.tree.root_node)
        for function in self.function_definitions.values():
            function["calls"] = list(function["calls"])

        # json.dump(self.function_definitions, open("functions.json", "w"), indent=4)
        return
    
    def analyze_function_declarations(self) -> None:
        def traverse(node: Node):
            if node.type == "declaration":
                for child in node.children:
                    if child.type == "function_declarator":
                        exception_flag = False
                        function_name = None
                        function_parameters = []
                        for grandchild in child.children:
                            if grandchild.type == "identifier":
                                function_name = grandchild.text.decode("utf-8")
                                continue
                            if grandchild.type == "parameter_list":
                                function_parameters = self.analyze_function_parameters(grandchild)
                                if function_parameters is False:
                                    exception_flag = True
                                continue
                        if function_name:
                            function_start = node.start_point
                            function_end = node.end_point
                            self.function_declarations[function_name] = {
                                "start": function_start,
                                "end": function_end,
                                "parameters": function_parameters,
                                "exception_flag": exception_flag,
                            }
                        break

            for child in node.children:
                traverse(child)

        if self.tree and self.tree.root_node:
            traverse(self.tree.root_node)
        # json.dump(self.function_declarations, open("function_declarations.json", "w"), indent=4)
        return
    
    def classify_function_calls(self) -> None:
        defined_functions = set(self.function_definitions.keys())
        declared_functions = set(self.function_declarations.keys())
        self.classified_functions = {"defined": set(), "declared_only": set(),"undefined": set()}

        for function, details in self.function_definitions.items():
            self.classified_functions["defined"].add(function)
            for called_function in details["calls"]:
                if called_function in defined_functions:
                    self.classified_functions["defined"].add(called_function)
                elif called_function in declared_functions:
                    self.classified_functions["declared_only"].add(called_function)
                else:
                    self.classified_functions["undefined"].add(called_function)

        for key in self.classified_functions:
            self.classified_functions[key] = list(self.classified_functions[key])
        # json.dump(self.classified_functions, open("classified_calls.json", "w"), indent=4)
        return
    
    def analyze_global_and_external_variables(self) -> None:
        identifierget = IdentifierGet()
        typeget = TypeGet()
        _declarator = IdentifierGet.get_subtype_of_declarator(identifierget)
        type_specifier = TypeGet.get_subtype_of_type_specifier(typeget)

        def get_type(node: Node):
            if node.type in type_specifier:
                function_name = node.type
                return getattr(TypeGet, function_name)(node)
            else:
                print("Error: out of exception when get type of global_or_external_variables. node type: " + node.type + " :: " + node.text.decode("utf-8"))
                return False

        def get_identifier(node: Node):
            if node.type in _declarator:
                function_name = node.type
                return getattr(IdentifierGet, function_name)(node)
            else:
                print("Error: out of exception when get identifier of global_or_external_variables. node type: " + node.type + " :: " + node.text.decode("utf-8"))
                return False

        for node in self.tree.root_node.children:
            if node.type == "declaration":
                if any(child.type == "function_declarator" for child in node.children):
                    # Skip function declarations
                    continue
                is_extern = False
                storage_class_specifier = None
                variable_declarator = None
                variable_type = None
                variable_start = node.start_point
                variable_end = node.end_point
                type_position = 0
                if node.named_children[0].type in ["type_qualifier", "storage_class_specifier"]:
                    storage_class_specifier = node.named_children[0].text.decode("utf-8")
                    type_position = 1
                    if storage_class_specifier == "extern":
                        is_extern = True
                variable_type = get_type(node.named_children[type_position])
                for var in node.named_children[type_position+1:]:
                    if var.type == "init_declarator":
                        variable_declarator = get_identifier(var.named_children[0])
                    else:
                        variable_declarator = get_identifier(var)

                    if is_extern:
                        if variable_type and variable_declarator:
                            variable_name = variable_declarator["identifier"]
                            if variable_name:
                                self.external_variables[variable_name] = {
                                    "declarator": variable_declarator,
                                    "type": variable_type,
                                    "start": variable_start,
                                    "end": variable_end,
                                }
                        else:
                            self.external_variables[str(node.start_point)] = {
                                "error": 'unknown',
                                "start": node.start_point,
                                "end": node.end_point,
                            }
                    else:
                        if variable_type and variable_declarator:
                            variable_name = variable_declarator["identifier"]
                            if variable_name:
                                self.global_variables[variable_name] = {
                                    "declarator": variable_declarator,
                                    "type": variable_type,
                                    "start": variable_start,
                                    "end": variable_end,
                                }
                        else:
                            self.global_variables[str(node.start_point)] = {
                                "error": 'unknown',
                                "start": node.start_point,
                                "end": node.end_point,
                            }

        for value in self.global_variables.values():
            if "type" in value.keys():
                if value["type"]["type_identifier"] in self.defined_types.keys() and \
                    value["type"]["type"] == self.defined_types[value["type"]["type_identifier"]]["type"]:
                    value["valid"] = True
                else:
                    value["valid"] = False

        json.dump(self.global_variables, open("global_variables.json", "w"), indent=4)
        json.dump(self.external_variables, open("external_variables.json", "w"), indent=4)
        return
    
    def analyze_function_return_type(self, return_type: Node) -> dict:
        typeget = TypeGet()
        type_specifier = typeget.get_subtype_of_type_specifier()

        if return_type.type in type_specifier:
            function_name = return_type.type
            return getattr(TypeGet, function_name)(return_type)
        else:
            print("Error: out of exception when get type of function_return_type. type: " + return_type.type + " :: " + return_type.text.decode("utf-8"))
            return False

    def analyze_function_parameters(self, parameter_list: Node) -> list:
        identifierget = IdentifierGet()
        typeget = TypeGet()
        _declarator = IdentifierGet.get_subtype_of_declarator(identifierget)
        _abstract_declarator = IdentifierGet.get_subtype_of_abstract_declarator(identifierget)
        type_specifier = TypeGet.get_subtype_of_type_specifier(typeget)
        
        def get_type(node: Node):
            if node.type in type_specifier:
                function_name = node.type
                return getattr(TypeGet, function_name)(node)
            else:
                print("Error: out of exception when get type of function_parameters. function: " + node.type + " :: " + node.text.decode("utf-8"))
                return False

        def get_identifier(node: Node):
            if node.type in _declarator:
                function_name = node.type
                return getattr(IdentifierGet, function_name)(node)
            elif node.type in _abstract_declarator:
                function_name = node.type
                return getattr(IdentifierGet, function_name)(node)
            else:
                print("Error: out of exception when get identifier of function_parameters. function: " + node.type + " :: " + node.text.decode("utf-8"))
                return False

        para_list = list()
        for parameter in parameter_list.named_children:
            if parameter.type == "parameter_declaration":
                if len(parameter.children) == 2:
                    para_list.append({"type": get_type(parameter.children[0]),
                                      "declarator": get_identifier(parameter.children[1])})
                elif len(parameter.children) == 1:
                    para_list.append({"type": get_type(parameter.children[0])})
            else:
                return False
            
        for para in para_list:
            if False in para.values():
                return False
        return para_list
    
    def analyze_local_variables(self, function_compound_statement: Node) -> dict:
        local_variables = {}
        identifierget = IdentifierGet()
        typeget = TypeGet()
        _declarator = IdentifierGet.get_subtype_of_declarator(identifierget)
        type_specifier = TypeGet.get_subtype_of_type_specifier(typeget)
        
        def get_type(node: Node):
            if node.type in type_specifier:
                function_name = node.type
                return getattr(TypeGet, function_name)(node)
            else:
                print("Error: out of exception when get type of local_variables. function: " + node.type + " :: " + node.text.decode("utf-8"))
                return False

        def get_identifier(node: Node):
            if node.type in _declarator:
                function_name = node.type
                return getattr(IdentifierGet, function_name)(node)
            else:
                print("Error: out of exception when get identifier of local_variables. function: " + node.type + " :: " + node.text.decode("utf-8"))
                return False

        def traverse(node: Node):
            if node.type == "declaration":
                variable_name = None
                type_position = 0
                if node.named_children[0].type in ["type_qualifier", "storage_class_specifier"]:
                    type_position = 1
                if len(node.named_children) >= 2:
                    variable_start = node.start_point
                    variable_end = node.end_point
                    variable_type = get_type(node.named_children[type_position])
                    for var in node.named_children[type_position+1:]:
                        if var.type == "init_declarator":
                            variable_declarator = get_identifier(var.named_children[0])
                        else:
                            variable_declarator = get_identifier(var)

                        if variable_type and variable_declarator:
                            variable_name = variable_declarator["identifier"]
                            if variable_name:
                                local_variables[variable_name] = {
                                    "declarator": variable_declarator,
                                    "type": variable_type,
                                    "start": variable_start,
                                    "end": variable_end,
                                }
                        else:
                            local_variables[str(node.start_point)] = {
                                "error": 'unknown',
                                "start": node.start_point,
                                "end": node.end_point,
                            }

            for child in node.children:
                traverse(child)

        traverse(function_compound_statement)
            
        # json.dump(local_variables, open("local_variables.json", "w"), indent=4)
        for value in local_variables.values():
            if "error" in value.keys():
                return False
        return local_variables
    
    def analyze_identifiers(self, root) -> list:
        identifiers = []

        def traverse(node: Node):
            if node.type == "identifier":
                identifier_name = node.text.decode("utf-8")
                identifiers.append(identifier_name)
            for child in node.children:
                traverse(child)

        traverse(root)

        return identifiers
    
    def analyze_struct_definitions(self) -> None:
        def traverse(node: Node):
            if node.type == "struct_specifier":
                struct_name = None
                is_definition = False
                for child in node.named_children:
                    if child.type == "type_identifier":
                        struct_name = child.text.decode("utf-8")
                        continue
                    if child.type == "field_declaration_list":
                        is_definition = True
                        continue
                if is_definition and struct_name:
                    struct_start = node.start_point
                    struct_end = node.end_point
                    self.defined_types[struct_name] = {
                        "name": struct_name,
                        "start": struct_start,
                        "end": struct_end,
                        "type": "struct",
                    }

            for child in node.children:
                traverse(child)

        if self.tree and self.tree.root_node:
            traverse(self.tree.root_node)
        return
    
    def analyze_function_transferability(self) -> None:
        def check_parameter(function_details):
            for parameter in function_details["parameters"]:
                if parameter["type"]["type"] in ["struct", "class", "union"]:
                    if parameter["type"]["type_identifier"] in self.defined_types.keys() and \
                        parameter["type"]["type"] == self.defined_types[parameter["type"]["type_identifier"]]["type"]:
                        continue
                    else:
                        return False
                else:
                    pass
            return True

        def check_local_variables(function_details):
            for variable_name, variable_details in function_details["local_variables"].items():
                try:
                    if variable_details["type"]["type"] in ["struct", "class", "union"]:
                        if variable_details["type"]["type_identifier"] in self.defined_types.keys() and \
                            variable_details["type"]["type"] == self.defined_types[variable_details["type"]["type_identifier"]]["type"]:
                            continue
                        else:
                            return False
                    else:
                        pass
                except:
                    pass
            return True

        def check_identifiers(function_details):
            parameter_list = []
            for x in function_details["parameters"]:
                try:
                    parameter_list.append(x["declarator"]["identifier"])
                except KeyError:
                    pass
            for identifier in function_details["identifiers"]:
                if identifier in function_details["local_variables"].keys():
                    continue
                elif identifier in parameter_list:
                    continue
                elif identifier in function_details["calls"]:
                    continue
                elif identifier in self.global_variables.keys() and \
                    self.global_variables[identifier]["valid"] is True:
                    continue
                elif identifier in self.defined_types.keys() and \
                    self.defined_types[identifier]["type"] == "macro":
                    continue
                else:
                    # print(f"identifier {identifier} unknown in function {function_details['start']}")
                    return False
            return True

        def check_function_call(function_details):
            for function_call in function_details["calls"]:
                if function_call == function_details["name"]:
                    continue
                if function_call not in self.function_definitions.keys():
                    return False
                if "transferability" not in self.function_definitions[function_call].keys():
                    do_analyze(self.function_definitions[function_call])
                if self.function_definitions[function_call]["transferability"] is False:
                    return False
                else:
                    continue
            return True
        
        def check_return_type(function_details):
            return_type = function_details["return_type"]
            if return_type["type"] in ["struct", "class", "union"]:
                if return_type["type_identifier"] in self.defined_types.keys() and \
                    return_type["type"] == self.defined_types[return_type["type_identifier"]]["type"]:
                    return True
                else:
                    return False
            else:
                return True

        def do_analyze(function_details):
            if function_details["exception_flag"] is True:
                function_details["transferability"] = False
                return
            if "transferability" in function_details.keys():
                return
            else:
                function_details["transferability_detail"] = [check_parameter(function_details),
                                                              check_local_variables(function_details),
                                                              check_identifiers(function_details),
                                                              check_function_call(function_details),
                                                              check_return_type(function_details)]
                function_details["transferability"] = all(function_details["transferability_detail"])
            return

        for function_details in self.function_definitions.values():
            do_analyze(function_details)
        json.dump(self.function_definitions, open("functions.json", "w"), indent=4)
        return